import tensorflow as tf
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

# 读取数据集
print("Loading dataset...")
data = pd.read_csv('sitting_posture_dataset.csv')

# 分离特征和标签
X = data.drop('label', axis=1).values
y = data['label'].values

# 数据分割
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# 数据标准化
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

# 保存数据标准化参数（用于Android端）
scaler_params = {
    'mean': scaler.mean_.tolist(),
    'scale': scaler.scale_.tolist()
}
import json
with open('scaler_params.json', 'w') as f:
    json.dump(scaler_params, f)

# 创建模型
print("Creating model...")
model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu', input_shape=(6,)),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(32, activation='relu'),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(16, activation='relu'),
    tf.keras.layers.Dense(1, activation='sigmoid')
])

# 编译模型
model.compile(
    optimizer='adam',
    loss='binary_crossentropy',
    metrics=['accuracy']
)

# 训练模型
print("Training model...")
history = model.fit(
    X_train_scaled, y_train,
    epochs=100,
    batch_size=32,
    validation_split=0.2,
    verbose=1
)

# 评估模型
print("\nEvaluating model...")
test_loss, test_accuracy = model.evaluate(X_test_scaled, y_test, verbose=0)
print(f"Test accuracy: {test_accuracy:.4f}")

# 保存完整模型
print("\nSaving full model...")
model.save('sitting_posture_model.h5')

# 转换为TFLite格式
print("\nConverting to TFLite...")
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

# 保存TFLite模型
with open('sitting_posture_model.tflite', 'wb') as f:
    f.write(tflite_model)

# 测试TFLite模型
print("\nTesting TFLite model...")
interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 准备测试数据
test_data = X_test_scaled[0:1]
interpreter.set_tensor(input_details[0]['index'], test_data.astype(np.float32))

# 运行推理
interpreter.invoke()
tflite_output = interpreter.get_tensor(output_details[0]['index'])

print("\nModel files created:")
print("1. sitting_posture_model.tflite - TFLite model for Android")
print("2. scaler_params.json - Standardization parameters for preprocessing")
print("3. sitting_posture_model.h5 - Full TensorFlow model (for reference)")

print("\nDone! The TFLite model is ready for Android implementation.") 